// 引入所需的包和类
package com.sprint.questai.module.llm;

import com.sprint.questai.factory.ModelFactory;
import com.sprint.questai.model.dto.request.ChatRequest;
import com.sprint.questai.model.entity.UserCaseInfo;
import com.sprint.questai.model.enums.ChatModelEnums;
import com.sprint.questai.model.enums.NameEnums;
import com.sprint.questai.module.prompt.Prompt;
import com.sprint.questai.module.store.ChatMemoryStore.ChatMemoryStoreInMemory;
import com.sprint.questai.service.Lawyer;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.router.LanguageModelQueryRouter;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static com.sprint.questai.module.prompt.ExtractPrompt.updateUserInfo;

/**
 * @Author: mayixiang
 * @Date: 2024-05-25
 */

@Service  // 标记为Spring服务组件
@Slf4j  // Lombok库提供的日志功能，自动创建log对象
public class AiChatService {
    @Autowired  // Spring自动注入相应的组件
    ModelFactory modelFactory;  // 模型工厂，用于获取不同的模型实例
    @Autowired
    ChatMemoryStoreInMemory chatMemoryStoreInMemory;  // 内存中的聊天历史存储
    PromptTemplate injectPrompt=PromptTemplate.from("{{userMessage}}\n\n" +
            "可以参考下面的信息来回答，下面信息包括法条，司法解释和具体案例，对于法条和重要的司法解释，应当相信并可以列出，" +
            "对于具体的案例，应该作为参考，参考法官判案的逻辑，此类案件可能需要的证据等\n{{contents}}");
    PromptTemplate template = PromptTemplate.from(
            "根据用户查询，从以下选项中较合适和问题有关的一个或者几个数据来源来检索相关信息,：\n{{options}}\n非常重要的是，你的答案只能由一个数字或多个用逗号分隔的数字组成，不能有其他任何内容！\n用户查询：{{query}}"
    );
    ChatLanguageModel languageModel;
    Map<String, UserCaseInfo> userCaseInfoMap=new HashMap<>();  // 用户案例信息的映射
    // 定义一个公共方法，用于处理聊天功能，输入参数为ChatRequest对象
    public String aiChatWithLocalKeyAndMemory(ChatRequest request){
        String model = request.getModel();  // 获取请求中指定的模型名
        String content = request.getContent();  // 获取聊天内容
        String memoryId = request.getUserId() + request.getConversationId(); // 根据用户ID和对话ID生成唯一的内存ID
        boolean useRag=request.isUseRag();
        List<ChatMessage> historyMessage = chatMemoryStoreInMemory.getMessages(memoryId);  // 获取当前对话ID的历史消息列表
        // 向聊天内存中追加用户消息
        ChatLanguageModel chatModel = ChatModelEnums.findModel(model);
        languageModel=chatModel;
        chatMemoryStoreInMemory.appendMessages(memoryId, Arrays.asList(new UserMessage(content)));
        if(content.contains(Prompt.extraInfoCmd)){
            UserCaseInfo userCaseInfo = updateUserCaseInfo(memoryId, content);
            return "目前您的信息如下"+userCaseInfo.toString();
        }else if(content.contains(Prompt.helpCmd)){
            return Prompt.helpPrompt;
        }else if(content.contains(Prompt.analyzeCmd)){
            UserCaseInfo userCaseInfo = userCaseInfoMap.get(memoryId);
            String s = analyzeCase(chatModel, historyMessage, userCaseInfo);
            chatMemoryStoreInMemory.appendMessages(memoryId, Arrays.asList(new AiMessage(s)));
            return s;
        }else if(content.contains(Prompt.endCmd)){
            return endChat(memoryId);
        }
        // 通过枚举查找对应的聊天模型
        return normalChat(memoryId);
    }


    public String  normalChat(String memoryId){
        Response<AiMessage> response = null;  // 初始化响应对象
        List<ChatMessage> historyMessage = chatMemoryStoreInMemory.getMessages(memoryId);  // 获取当前对话ID的历史消息列表
        // 如果找到了聊天模型，就使用该模型生成回复
        if (languageModel != null) {
            response = languageModel.generate(historyMessage);
            log.info("prompt: " + historyMessage);  // 记录日志，输出历史消息
        }else{
            throw new RuntimeException("model not found: " );  // 如果没有找到模型，记录错误日志
        }
        // 如果生成了响应，并且终止原因为STOP，那么将AI的回复也追加到聊天内存中
        if (response != null && response.finishReason().equals(FinishReason.STOP)) {
            chatMemoryStoreInMemory.appendMessages(memoryId, Arrays.asList(new AiMessage(response.content().text())));
        }

        // 返回AI的回复内容
        return response.content().text();
    }
    public UserCaseInfo updateUserCaseInfo(String memoryId,String info){
        UserCaseInfo userCaseInfo = userCaseInfoMap.get(memoryId);
        userCaseInfo = updateUserInfo(userCaseInfo, info);
        userCaseInfoMap.put(memoryId, userCaseInfo);
        return userCaseInfo;
    }
    public String endChat(String memoryId){

        EmbeddingStore caseEmbeddingStore = modelFactory.createEmbeddingStore(
                NameEnums.CASE_RECORD_NAME.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingStore  explainStore = modelFactory.createEmbeddingStore(NameEnums.EXPLAIN_NAME.toString(), NameEnums.DEFAULT_EMBEDDING_MODEL.toString(), NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingModel defaultEmbeddingModel = modelFactory.createEmbeddingModel(NameEnums.BGE_SMALL_ZH.toString(),"");
        EmbeddingStore lawStore = modelFactory.createEmbeddingStore(
                NameEnums.LAW_NAME.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString());
        // given
        ContentRetriever casecontentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(caseEmbeddingStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(3)
                .build();
        ContentRetriever judgeReasonContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(explainStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(5)
                .build();
        ContentRetriever lawContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(lawStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(5)
                .build();

        Map<ContentRetriever, String> retrieverToDescription = new HashMap<>();
        retrieverToDescription.put(casecontentRetriever, NameEnums.CASE_DESCRIPTION.toString());
        retrieverToDescription.put(judgeReasonContentRetriever, NameEnums.EXPLAIN_DESCRIPTION.toString());
        retrieverToDescription.put(lawContentRetriever, NameEnums.LAW_DESCRIPTION.toString());
        ContentAggregator aggregator = new DefaultContentAggregator();
        Lawyer lawyer = AiServices.builder(Lawyer.class)
                .retrievalAugmentor(DefaultRetrievalAugmentor.
                                builder()
                                .contentAggregator(aggregator)
                                .queryRouter(
                                        LanguageModelQueryRouter.builder().chatLanguageModel(languageModel)
                                                .promptTemplate(template)
                                                .retrieverToDescription(retrieverToDescription)
                                                .build()
                                )
                                .contentInjector(new DefaultContentInjector(injectPrompt))
//                        .queryRouter(new LanguageModelQueryRouter(model,retrieverToDescription))
                                .build()
                )
                .chatLanguageModel(modelFactory.createChatLanguageModel(NameEnums.GLM4E.toString()))
                .build();
        List<ChatMessage> messages = chatMemoryStoreInMemory.getMessages(memoryId);
        UserCaseInfo userCaseInfo=userCaseInfoMap.get(memoryId);
        if(userCaseInfo==null)userCaseInfo=new UserCaseInfo();
        String baogao = lawyer.endAndAnalyze( messages,userCaseInfo);
        log.info(baogao);
        return baogao;
    }

    public String analyzeCase(ChatLanguageModel model, String query,UserCaseInfo info){
        //将query转换为ChatMessage对象
        List<ChatMessage> messages = Arrays.asList(new UserMessage(query));
        return analyzeCase(model, messages,info);
    }
    public String analyzeCase(ChatLanguageModel model, List<ChatMessage> messages,UserCaseInfo info){
         PromptTemplate template = PromptTemplate.from(
                 "根据用户查询，从以下选项中较合适和问题有关的一个或者几个数据来源来检索相关信息,：\n{{options}}\n非常重要的是，你的答案只能由一个数字或多个用逗号分隔的数字组成，不能有其他任何内容！\n用户查询：{{query}}"
         );
        EmbeddingStore caseEmbeddingStore = modelFactory.createEmbeddingStore(
                NameEnums.CASE_RECORD_NAME.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingStore  explainStore = modelFactory.createEmbeddingStore(NameEnums.EXPLAIN_NAME.toString(), NameEnums.DEFAULT_EMBEDDING_MODEL.toString(), NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingModel defaultEmbeddingModel = modelFactory.createEmbeddingModel(NameEnums.BGE_SMALL_ZH.toString(),"");
        EmbeddingStore lawStore = modelFactory.createEmbeddingStore(
                NameEnums.LAW_NAME.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString());
        // given
        ContentRetriever casecontentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(caseEmbeddingStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(3)
                .build();
        ContentRetriever judgeReasonContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(explainStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(5)
                .build();
        ContentRetriever lawContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(lawStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(5)
                .build();

        Map<ContentRetriever, String> retrieverToDescription = new HashMap<>();
        retrieverToDescription.put(casecontentRetriever, NameEnums.CASE_DESCRIPTION.toString());
        retrieverToDescription.put(judgeReasonContentRetriever, NameEnums.EXPLAIN_DESCRIPTION.toString());
        retrieverToDescription.put(lawContentRetriever, NameEnums.LAW_DESCRIPTION.toString());
        ContentAggregator aggregator = new DefaultContentAggregator();
        Lawyer lawyer = AiServices.builder(Lawyer.class)
                .retrievalAugmentor(DefaultRetrievalAugmentor.
                        builder()
                        .contentAggregator(aggregator)
                        .queryRouter(
                                LanguageModelQueryRouter.builder().chatLanguageModel(languageModel)
                                        .promptTemplate(template)
                                        .retrieverToDescription(retrieverToDescription)
                                        .build()
                        )
                                .contentInjector(new DefaultContentInjector(injectPrompt))
//                        .queryRouter(new LanguageModelQueryRouter(model,retrieverToDescription))
                        .build()
                )
                .chatLanguageModel(modelFactory.createChatLanguageModel(NameEnums.GLM4E.toString()))
                .build();
        String keyFactors = lawyer.findKeyFactors(messages);

        return keyFactors;
    }
    public String normalChatWithRag(String query,String memoryId){
        EmbeddingStore caseEmbeddingStore = modelFactory.createEmbeddingStore(
                NameEnums.CASE_RECORD_NAME.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingStore  explainStore = modelFactory.createEmbeddingStore(NameEnums.EXPLAIN_NAME.toString(), NameEnums.DEFAULT_EMBEDDING_MODEL.toString(), NameEnums.ELASTIC_SEARCH_STORE.toString());
        EmbeddingModel defaultEmbeddingModel = modelFactory.createEmbeddingModel(NameEnums.BGE_SMALL_ZH.toString(),"");
        // given
        ContentRetriever casecontentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(caseEmbeddingStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(3)
                .build();
        ContentRetriever judgeReasonContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(explainStore)
                .embeddingModel(defaultEmbeddingModel)
                .maxResults(3)
                .build();
        Map<ContentRetriever, String> retrieverToDescription = new HashMap<>();
        retrieverToDescription.put(casecontentRetriever, "记录案件判决.包括案件事实,判决结果,判决引用法条等");
        retrieverToDescription.put(judgeReasonContentRetriever, "记录法律条文,包括刑法,公司法等");
        ContentAggregator aggregator = new DefaultContentAggregator();
        Lawyer lawyer = AiServices.builder(Lawyer.class)
                .retrievalAugmentor(DefaultRetrievalAugmentor.
                        builder()
                        .contentAggregator(aggregator)
//                        .queryRouter(
//                                LanguageModelQueryRouter.builder().chatLanguageModel(languageModel)
//                                        .promptTemplate()
//                        )
                        .queryRouter(new LanguageModelQueryRouter(languageModel,retrieverToDescription))
                        .build()
                )
                .chatLanguageModel(modelFactory.createChatLanguageModel(NameEnums.GLM4E.toString()))
                .build();
        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder().queryRouter(
                new LanguageModelQueryRouter(languageModel, retrieverToDescription)).build();
        UserMessage augment = retrievalAugmentor.augment(UserMessage.from(query), null);
        chatMemoryStoreInMemory.appendMessages(memoryId, Arrays.asList(augment));
        Response<AiMessage> aiMessageResponse = languageModel.generate(chatMemoryStoreInMemory.getMessages(memoryId));
        chatMemoryStoreInMemory.appendMessages(memoryId, Arrays.asList(aiMessageResponse.content()));
        return aiMessageResponse.content().text();

    }
}
